{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q9KfUf1BI6Kl"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "WvqLCVQ6I58i"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZoFRICPTNUca"
      },
      "source": [
        "# SavedModel Migration\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/migrate/saved_model\">\n",
        "    <img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />\n",
        "    View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/migrate/saved_model.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />\n",
        "    Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/migrate/saved_model.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />\n",
        "    View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/migrate/saved_model.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nGyIb7MYJfaM"
      },
      "source": [
        "Once you have migrated your model from graphs and sessions to `tf.function`, `tf.Module`, and `tf.keras.Model`, it is time to migrate the saving and loading code. This is specifically a guide about migrating from TF1 to TF2, [see here a more general guide about the TF2 SavedModel API](../../guide/saved_model.ipynb).\n",
        "\n",
        "This is a quick overview of the API changes from TF1 to TF2.\n",
        "\n",
        "||TF1|Migration to TF2|\n",
        "| ----| ----|---|\n",
        "|**Saving**|`tf.compat.v1.saved_model.Builder`<br>`tf.compat.v1.saved_model.simple_save`|`tf.saved_model.save`<br>`tf.keras.models.save_model`\n",
        "|**Loading**|`tf.compat.v1.saved_model.load`|`tf.saved_model.load`\n",
        "|**Signatures**: a set of input <br>and output tensors that <br>can be used to run the<br> model|These are generated using the `*.signature_def` utils<br>(e.g. `tf.compat.v1.saved_model.predict_signature_def`)|Write a `tf.function`, and export it using the<br> `signatures` argument in <br>`tf.saved_model.save`.\n",
        "**Classify and regress**:<br>special types of signatures|Generated with `classification_signature_def`, <br> `regresssion_signature_def`, and certain Estimator exports.|These two signatures types have been removed <br>from TF2.If the serving library requires these <br>method names, use the [`MethodNameUpdater`](https://www.tensorflow.org/api_docs/python/tf/compat/v1/saved_model/signature_def_utils/MethodNameUpdater).\n",
        "\n",
        "For a more in-depth explanation of the mapping see the *Changes from TF1 to TF2* section\n",
        "\n",
        "Below are code examples of writing saving and loading code in TF1 and TF2."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r5mR2xsNAGsB"
      },
      "source": [
        "## Code Sample Set up\n",
        "\n",
        "The code examples below show how to export and load the same dummy TensorFlow model (`add_two`) to SavedModel using the TF1 and TF2 APIs.\n",
        "\n",
        "Run the below code to set up the imports and utils functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B94QZyy-kOGQ"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow.compat.v1 as tf1\n",
        "import shutil\n",
        "\n",
        "def remove_dir(path):\n",
        "  try:\n",
        "    shutil.rmtree(path)\n",
        "  except:\n",
        "    pass\n",
        "    \n",
        "def add_two(input):\n",
        "  return input + 2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZNVpH5tOCgd9"
      },
      "source": [
        "## Building TF1 SavedModels\n",
        "\n",
        "\n",
        "The TF1 APIs, `tf.compat.v1.saved_model.Builder`, `tf.compat.v1.saved_model.simple_save`, and `tf.estimator.Estimator.export_saved_model` export the TensorFlow graph and session."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "THRLul5ijmTE"
      },
      "source": [
        "### Builder\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dcZDQaI8jl3h"
      },
      "outputs": [],
      "source": [
        "remove_dir(\"saved-model-builder\")\n",
        "\n",
        "with tf.Graph().as_default() as g:\n",
        "  with tf1.Session() as sess:\n",
        "    input = tf1.placeholder(tf.float32, shape=[])\n",
        "    output = add_two(input)\n",
        "    print(\"add two output: \", sess.run(output, {input: 3.}))\n",
        "\n",
        "    # Save with SavedModelBuilder\n",
        "    builder = tf1.saved_model.Builder('saved-model-builder')\n",
        "    sig_def = tf1.saved_model.predict_signature_def(\n",
        "        inputs={'input': input}, \n",
        "        outputs={'output': output})\n",
        "    builder.add_meta_graph_and_variables(\n",
        "        sess, tags=[\"serve\"], signature_def_map={\n",
        "            tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: sig_def\n",
        "    })\n",
        "    builder.save()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PwtC27VFlwCa"
      },
      "outputs": [],
      "source": [
        "!saved_model_cli run --dir simple-save --tag_set serve \\\n",
        " --signature_def serving_default --input_exprs input=10"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gnBDNTxKG_vR"
      },
      "source": [
        "### Simple Save"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jtMxe2rjHSq9"
      },
      "outputs": [],
      "source": [
        "remove_dir(\"simple-save\")\n",
        "\n",
        "with tf.Graph().as_default() as g:\n",
        "  with tf1.Session() as sess:\n",
        "    input = tf1.placeholder(tf.float32, shape=[])\n",
        "    output = add_two(input)\n",
        "    print(\"add two output: \", sess.run(output, {input: 3.}))\n",
        "\n",
        "    tf1.saved_model.simple_save(\n",
        "        sess, 'simple-save',\n",
        "        inputs={'input': input}, \n",
        "        outputs={'output': output})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AdnqemvIHb2P"
      },
      "outputs": [],
      "source": [
        "!saved_model_cli run --dir simple-save --tag_set serve \\\n",
        " --signature_def serving_default --input_exprs input=10"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r0BNzzAHjnkp"
      },
      "source": [
        "### Estimator export\n",
        "\n",
        "In the definition of the Estimator `model_fn`, you can define signatures in your model by returning `export_outputs` in the `EstimatorSpec`. There are different types of outputs:\n",
        "* `tf.estimator.export.ClassificationOutput`\n",
        "* `tf.estimator.export.RegressionOutput`\n",
        "* `tf.estimator.export.PredictOutput`\n",
        "\n",
        "These will produce `classify`, `regress`, and `predict` signature types. \n",
        "\n",
        "When the estimator is exported with `tf.estimator.Estimator.export_saved_model`, these signatures will be saved with the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3nQ5Stnxjhfs"
      },
      "outputs": [],
      "source": [
        "def model_fn(features, labels, mode):\n",
        "  output = add_two(features['input'])\n",
        "  step = tf1.train.get_global_step()\n",
        "  return tf.estimator.EstimatorSpec(\n",
        "      mode,\n",
        "      predictions=output,\n",
        "      train_op=step.assign_add(1),\n",
        "      loss=tf.constant(0.),\n",
        "      export_outputs={\n",
        "          tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: \\\n",
        "          tf.estimator.export.PredictOutput({'output': output})})\n",
        "est = tf.estimator.Estimator(model_fn, 'estimator-checkpoints')\n",
        "\n",
        "# Train for one step to create a checkpoint\n",
        "def train_fn():\n",
        "  return tf.data.Dataset.from_tensors({'input': 3.})\n",
        "est.train(train_fn, steps=1)\n",
        "\n",
        "# This util function `build_raw_serving...` takes in raw tensor features\n",
        "# and builds an \"input serving receiver function\", which creates placeholder\n",
        "# inputs to the model.\n",
        "serving_input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn(\n",
        "    {'input': tf.constant(3.)})  # Pass in a dummy input batch\n",
        "estimator_path = est.export_saved_model('exported-estimator', serving_input_fn)\n",
        "\n",
        "# Estimator's export_saved_model creates a timestamped directory. Move this\n",
        "# to a set path so it can be inspected with saved_model_cli in the cell below:\n",
        "!rm -rf estimator-model\n",
        "import shutil\n",
        "shutil.move(estimator_path, 'estimator-model')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8_gD2gkE7CMu"
      },
      "outputs": [],
      "source": [
        "!saved_model_cli run --dir estimator-model --tag_set serve \\\n",
        " --signature_def serving_default --input_exprs input=[10]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DyBvrNQgIhIo"
      },
      "source": [
        "## Building TF2 SavedModels"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BZmFH-eIjqjB"
      },
      "source": [
        "### TF2 Saving\n",
        "\n",
        "To export your model in TF2, you must define a `tf.Module` or `tf.keras.Model` to hold all of your model's variables and functions. Then, call `tf.saved_model.save` to create a SavedModel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_j-PwgP_jrgw"
      },
      "outputs": [],
      "source": [
        "class MyModel(tf.Module):\n",
        "  @tf.function\n",
        "  def __call__(self, input):\n",
        "    return add_two(input)\n",
        "\n",
        "model = MyModel()\n",
        "\n",
        "@tf.function\n",
        "def serving_default(input):\n",
        "  return {'output': model(input)}\n",
        "\n",
        "signature_function = serving_default.get_concrete_function(\n",
        "    tf.TensorSpec(shape=[], dtype=tf.float32))\n",
        "tf.saved_model.save(\n",
        "    model, 'tf2-save', signatures={\n",
        "        tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_function})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "slvU4vZN756F"
      },
      "outputs": [],
      "source": [
        "!saved_model_cli run --dir tf2-save --tag_set serve \\\n",
        " --signature_def serving_default --input_exprs input=10"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UYpSfbBJjr33"
      },
      "source": [
        "### Keras SavedModel (TF2)\n",
        "\n",
        "The Keras `tf` format exports a SavedModel from a `tf.keras.Model`. SavedModels exported with Keras can be reloaded using any of the loading APIs, in addition to the Keras loading function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mMcjhzyRjvp6"
      },
      "outputs": [],
      "source": [
        "inp = tf.keras.Input(3)\n",
        "out = add_two(inp)\n",
        "model = tf.keras.Model(inputs=inp, outputs=out)\n",
        "\n",
        "@tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.float32)])\n",
        "def serving_default(input):\n",
        "  return {'output': model(input)}\n",
        "\n",
        "model.save('keras-model', save_format='tf', signatures={\n",
        "        tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: serving_default})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4P93WP5R7-VT"
      },
      "outputs": [],
      "source": [
        "!saved_model_cli run --dir keras-model --tag_set serve \\\n",
        " --signature_def serving_default --input_exprs input=10"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SEKe9rGgoGCw"
      },
      "source": [
        "## Loading SavedModels\n",
        "\n",
        "SavedModels saved with any of the above APIs can be loaded with either TF1 or TF2. \n",
        "\n",
        "TF1 SavedModel can generally be used for inference when loaded into TF2, but training (generating gradients) is only possible if the SavedModel contains *resource variables*. You can check the dtype of the variables. If the variable dtype contains \"_ref\" then it is a reference variable.\n",
        "\n",
        "TF2 SavedModel can be loaded and executed from TF1 as long as the SavedModel is saved with signatures.\n",
        "\n",
        "The sections below contain code samples showing how to load the SavedModels saved in the previous sections, and call the exported signature."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hLztK_0YoTEP"
      },
      "source": [
        "### TF1 Loading\n",
        "\n",
        "TF1 imports the SavedModel directly in to the current graph and session. You can call `session.run` on the tensor input and output names."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IMO0laj-m0p9"
      },
      "outputs": [],
      "source": [
        "def load_tf1(path, input):\n",
        "  print('Loading from', path)\n",
        "  with tf.Graph().as_default() as g:\n",
        "    with tf1.Session() as sess:\n",
        "      meta_graph = tf1.saved_model.load(sess, [\"serve\"], path)\n",
        "      sig_def = meta_graph.signature_def[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]\n",
        "      input_name = sig_def.inputs['input'].name\n",
        "      output_name = sig_def.outputs['output'].name\n",
        "      print('  Output with input', input, ': ', \n",
        "            sess.run(output_name, feed_dict={input_name: input}))\n",
        "\n",
        "load_tf1('saved-model-builder', 5.)\n",
        "load_tf1('simple-save', 5.)\n",
        "load_tf1('estimator-model', [5.])  # Estimator's input must be batched.\n",
        "load_tf1('tf2-save', 5.)\n",
        "load_tf1('keras-model', 5.)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FbR3sfvooVBN"
      },
      "source": [
        "### TF2 Loading\n",
        "\n",
        "In TF2, objects are loaded into a Python object that stores the variables and functions. This is compatible with models saved from TF1. See `tf.saved_model.load` for details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OA52ezWV_KgL"
      },
      "outputs": [],
      "source": [
        "def load_tf2(path, input):\n",
        "  print('Loading from', path)\n",
        "  loaded = tf.saved_model.load(path)\n",
        "  out = loaded.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY](\n",
        "      tf.constant(input))['output']\n",
        "  print('  Output with input', input, ': ', out)\n",
        "\n",
        "load_tf2('saved-model-builder', 5.)\n",
        "load_tf2('simple-save', 5.)\n",
        "load_tf2('estimator-model', [5.])  # Estimator's input must be batched.\n",
        "load_tf2('tf2-save', 5.)\n",
        "load_tf2('keras-model', 5.)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gz3VFn5aAfmK"
      },
      "source": [
        "Models saved with the TF2 API can also access tf.functions and variables that are attached to the model (instead of those exported as signatures). Example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IfMTp-TGAfOs"
      },
      "outputs": [],
      "source": [
        "loaded = tf.saved_model.load('tf2-save')\n",
        "print('restored __call__:', loaded.__call__)\n",
        "print('output with input 5.', loaded(5))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VMoErNKHoXEg"
      },
      "source": [
        "### Keras loading\n",
        "\n",
        "The Keras loading API allows you to reload a saved model back into a Keras Model object. Note that this only allows you to load SavedModels saved with Keras (`model.save` or `tf.keras.models.save_model`). Models saved with `tf.saved_model.save` can not be loaded.\n",
        "\n",
        "Load the previously saved Keras model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZFUAxK0YeIAe"
      },
      "outputs": [],
      "source": [
        "loaded_model = tf.keras.models.load_model('keras-model')\n",
        "loaded_model.predict_on_batch(tf.constant([1, 3, 4]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GZ5vGJ0IDorc"
      },
      "source": [
        "## Changes from TF1 to TF2\n",
        "This section lists out key terms from TF1, their TF2 equivalents, and what has changed.\n",
        "\n",
        "### SavedModel\n",
        "\n",
        "SavedModel is a format that stores a TensorFlow model. It contains signatures which are used by serving platforms to run the model.\n",
        "\n",
        "The file format itself has not changed significantly, so SavedModels can be loaded and served using either TF1 and TF2 APIs. \n",
        "\n",
        "**Differences between TF1 and TF2**\n",
        "\n",
        "The *serving* and *inference* use cases have not been updated, aside from API changes. The improvement in TF2 is the ability to *reuse* and *compose models* loaded from SavedModel.\n",
        "\n",
        "In TF2, the program is represented by objects like `tf.Variable`, `tf.Module`, or Keras models and layers. There are no more global variables that have values stored in a session, and the graph now exists in different `tf.functions`. Consequently, during export, SavedModel saves each component and function graphs separately.\n",
        "\n",
        "When you write a TensorFlow program with the TF2 Python API, you must build an object to manage the variables, functions, and other resources. Generally, this is accomplished by using the Keras API, but you can also build the object by creating or subclassing `tf.Module`.\n",
        "\n",
        "Keras models and `tf.Module` automatically track variables and functions attached to them. SavedModel saves these connections between modules, variables and functions so that they can be restored when loading.\n",
        "\n",
        "\n",
        "### Signatures\n",
        "\n",
        "Signatures are the endpoints of a SavedModel -- they tell the user how to run the model and what inputs are needed. \n",
        "\n",
        "In TF1, signatures are created by listing the input and output tensors. In TF2, signatures are generated by passing in *concrete functions*.\n",
        "\n",
        "To read more about TensorFlow functions, [see this guide](../guide/intro_to_graphs). In short, a concrete function is generated from a `tf.function`:\n",
        "```\n",
        "# Option 1: specify an input signature\n",
        "@tf.function(input_signature=[...])\n",
        "def fn(...):\n",
        "  ... \n",
        "  return outputs\n",
        "\n",
        "tf.saved_model.save(model, path, signatures={\n",
        "    'name': fn\n",
        "})\n",
        "```\n",
        "```\n",
        "# Option 2: call get_concrete_function\n",
        "@tf.function\n",
        "def fn(...):\n",
        "  ...\n",
        "  return outputs\n",
        "\n",
        "tf.saved_model.save(model, path, signatures={\n",
        "    'name': fn.get_concrete_function(...)\n",
        "})\n",
        "```\n",
        "\n",
        "### `session.run`\n",
        "\n",
        "In TF1, you could call `session.run` with the imported graph as long as you already know the tensor names. This allows you to retrieve the restored variable values, or run parts of the model that were not exported in the signatures.\n",
        "\n",
        "In TF2, you can directly access the variable (e.g. `loaded.dense_layer.kernel`), or call `tf.function`s attached the model object (e.g. `loaded.__call__`).\n",
        "\n",
        "Unlike TF1, there is no way to extract parts of a function and access intermediate values. You *must* export all of the needed functionality in the saved object.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b6NG9JvUwJxn"
      },
      "source": [
        "## Serving Migration notes\n",
        "\n",
        "SavedModel was originally created to work with [TensorFlow Serving](https://www.tensorflow.org/tfx/guide/serving). This platform offers different types of prediction requests: classify, regress, and predict.\n",
        "\n",
        "The TF1 API allows you to create these types of signatures with the utils: \n",
        "* `tf.compat.v1.saved_model.classification_signature_def`\n",
        "* `tf.compat.v1.saved_model.regression_signature_def`\n",
        "* `tf.compat.v1.saved_model.predict_signature_def`\n",
        "\n",
        "Classify and regress restrict the inputs and outputs, so the inputs must be a `tf.Example`, and the outputs must be `classes` or `scores` or `prediction`. Meanwhile, the predict signature has no restrictions.\n",
        "\n",
        "SavedModels exported with the **TF2** API are compatible with TensorFlow serving, but will only contain `predict` signatures. The `classify` and `regress` signatures have ben removed. \n",
        "\n",
        "If you require the use of the `classify` and `regress` signatures, you may modify the exported SavedModel using `tf.compat.v1.saved_model.signature_def_utils.MethodNameUpdater`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MttvGjoDEfnR"
      },
      "source": [
        "## Other Saving Migration guides\n",
        "\n",
        "If you are using TF Hub, then here are some guides that you may fine useful:\n",
        "* https://www.tensorflow.org/hub/model_compatibility\n",
        "* https://www.tensorflow.org/hub/migration_tf2\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "saved_model.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
